import base64
import hashlib
import json
from datetime import datetime
from datetime import timedelta

HEADER = {
    'type': 'jwt',
    'alg': 'sha256'
}


class JsonWebToken:
    def __init__(self, head=None):
        self.__head = head if head else HEADER
        self.__head = self.__json_dumps(self.__head)
        self.__head = self.__base64_encode(self.__head)

        self.__payload = None
        self.__token = None

    def jwt_payload_handler(self, user):
        """
        传入用户对象，返回payload, bytes对象
        """
        payload_dict = {
            'id': user.id,  # 用户唯一标识
            'username': user.username,  # 用户名
            'iat': datetime.now().timestamp(),  # 签发时间
            'exp': (datetime.now() + timedelta(days=7)).timestamp()  # 过期时间
        }

        return self.__json_dumps(payload_dict)

    def jwt_encode_handler(self, payload):
        """
        传入payload，返回token字符串
        """
        secret = self.__head + b"." + self.__base64_encode(payload)

        return secret + b'.' + hashlib.sha256(secret).hexdigest().encode()

    def jwt_decode_handler(self, token):
        """
        将token字符串传入，返回payload
        """
        if not isinstance(token, bytes):
            token = token.encode()

        self.__payload = token.split(b'.')[1]
        self.__token = token.split(b'.')[-1]
        return self.__payload

    def authenticated(self, payload):
        """
        校验登录是否过期
        """
        if not isinstance(payload, bytes):
            payload = payload.encode()

        payload_string = self.__base64_decode(payload)
        payload_dict = self.__json_loads(payload_string)
        exp = payload_dict.get('exp', None)
        if exp and exp > datetime.now().timestamp():
            return True
        return False

    @staticmethod
    def __base64_encode(string):
        """
        将string的base64编码返回
        """
        if not isinstance(string, bytes):
            string = string.encode()
        return base64.b64encode(string)

    @staticmethod
    def __base64_decode(string):
        """
        将string解码后返回
        """
        if not isinstance(string, bytes):
            string = string.encode()

        return base64.b64decode(string).decode()

    @staticmethod
    def __json_dumps(obj):
        """
        将obj序列化
        """
        return json.dumps(obj).encode()

    @staticmethod
    def __json_loads(string):
        """
        将string反序列化
        """
        return json.loads(string)


class User:

    def __init__(self, username, id=1):
        self.id = id
        self.username = username


if __name__ == '__main__':
    user = User('dyp')
    jwt = JsonWebToken()
    payload = jwt.jwt_payload_handler(user)
    token = jwt.jwt_encode_handler(payload)
    payload = jwt.jwt_decode_handler(token)
    user = jwt.jwt_user_from_payload_handler(payload)

    print(user.username)
